#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.api import VAR

# 读取标准化后的波动率面板
df = pd.read_csv('merged_volatility_panel.csv', index_col='trade_date', parse_dates=True)
df_std = (df - df.mean()) / df.std()

# 提取滚动窗口数据
window_size = 200
horizon = 10
window_data = df_std[-window_size:]

# 删除标准差为0的列（VAR不接受）
window_data = window_data.loc[:, window_data.std() > 1e-6]

# 重新拟合VAR模型（只保留可用列）
model = VAR(window_data)
result = model.fit(maxlags=1)

# 关键一步：获取实际参与建模的变量名
actual_assets = result.model.endog_names
if isinstance(actual_assets, str):
    actual_assets = [actual_assets]

#  强制提取建模数据子集（确保列名和顺序匹配）
window_data_final = window_data[actual_assets]

#  重新拟合VAR模型（强制与变量名一致）
model = VAR(window_data_final)
result = model.fit(maxlags=1)

# 方向性溢出函数
def compute_directional_spillovers(var_model, steps):
    fevd_raw = var_model.fevd(steps).decomp[-1]

    # 获取参与建模的资产名
    asset_names = var_model.model.endog_names

    # 裁剪成方阵
    fevd = fevd_raw[:10, :10]

    # 标准化
    fevd_norm = fevd / fevd.sum(axis=1, keepdims=True)
    np.fill_diagonal(fevd_norm, 0)
    
    
    #  直接保存为贡献矩阵 CSV（顺序与方向性指标一致）
    asset_names = var_model.model.endog_names
    spill_matrix = pd.DataFrame(fevd_norm, index=asset_names[:10], columns=asset_names[:10])
    spill_matrix.to_csv("dy_spillover_matrix.csv")
    print("已成功保存 N×N 贡献矩阵：dy_spillover_matrix.csv")
    

    to_ = np.array(fevd_norm.sum(axis=0)).flatten()
    from_ = np.array(fevd_norm.sum(axis=1)).flatten()
    net = to_ - from_

    # 打印长度检查
    print(f"资产数: {len(asset_names)}, from_: {len(from_)}, to_: {len(to_)}, net: {len(net)}")

    # 构建DataFrame
    min_len = min(len(asset_names), len(from_), len(to_), len(net))

    return pd.DataFrame({
        '资产': asset_names[:min_len],
        'From_Others（接收）': from_[:min_len],
        'To_Others（输出）': to_[:min_len],
        'Net_Spillover（净值）': net[:min_len]

    })



# 运行
dy_df = compute_directional_spillovers(result, steps=horizon)
dy_df.to_csv('dy_directional_spillover.csv', index=False)
print("溢出指标计算完成！前5行如下：")
print(dy_df.head())

# 可视化
plt.figure(figsize=(12, 5))
plt.bar(dy_df['资产'], dy_df['Net_Spillover（净值）'], color='tomato')
plt.axhline(0, linestyle='--', color='gray')
plt.title('Net asset risk contagion index(Net Spillover)')
plt.ylabel('Net infection intensity')
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig('net_spillover_barplot.png', dpi=300)
plt.show()


